package speech.tag;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.io.Reader;
import java.util.ArrayList;
import java.util.Hashtable;
import java.util.List;
import java.util.Map;
import java.util.Scanner;
import java.util.StringTokenizer;

import be.ac.ulg.montefiore.run.jahmm.ObservationInteger;
import be.ac.ulg.montefiore.run.jahmm.io.ObservationIntegerReader;
import be.ac.ulg.montefiore.run.jahmm.io.ObservationSequencesReader;

import edu.stanford.nlp.tagger.maxent.MaxentTagger;

import speech.nbc.NaiveBayesClassifier;
import speech.nbc.Tag;
import speech.utils.Parser;

public class SpeechActTagger {
	private ArrayList<String> dictionary;
	private Hashtable<String,String> testFiles;
	
	private HmmSpeechActTagger hmmSystem;
	private NaiveBayesClassifier nbc;
	
	private Map<String,Double> hmmThresholdModifiers;
	
	private final String TEXT_FOLDER = "text/";
	private final String POS_FOLDER = "pos/";
	private final String TRAIN_FOLDER = "train/";
	private final String TEST_FOLDER = "test/";
	private final String STEMMER_FREQ_FOLDER = "frequencies/";
	
	public static final String INPUT_FILE = "input.txt";
	private final String HMM_CONFIG_FILE = "config.txt";
	
	private final String POS_LOG_FILE = "POSlog.txt";
	private final String HMM_LOG_FILE = "HMMLog.txt";
	private final String NBC_LOG_FILE = "NBCLog.txt";
	private final String LOG_FILE = "Log.txt";
	
//	private final String TEST_HEADER = "Testing with keyword threshold 3 for all tags.\n\n";
	
	private final int HMM_DELTA = 2;
	private int HMM_THRESHOLD;
	
	public SpeechActTagger(int hmmT, Map<String, Double> hmmModifiers) {
		// Setup the system.
		System.out.println("SpeechActTagger v0.1\n\nSetting up the system:");
		this.HMM_THRESHOLD = hmmT;
    	hmmThresholdModifiers = hmmModifiers;
		setupSystem();
		// Initialize the HMM system.
		hmmSystem = new HmmSpeechActTagger(HMM_CONFIG_FILE, dictionary.size());		
    	nbc = new NaiveBayesClassifier(TEXT_FOLDER);
	}
	
	private void setupSystem() {
		Parser p = new Parser();
		
		System.out.println("Generating POS files...");
		processInputFile(INPUT_FILE, true);
		
		System.out.println("Generating dictionary...");
		dictionary = p.constructDictionaryFromPOSFiles("pos/");
		
		System.out.println("Encoding POS files, writing HMM config file...");
		processInputFile(INPUT_FILE, false);
	}
	
	public void processInputFile(String input, boolean posTagging) {
		PosTagger posTagger = null;
		Scanner s;
		PrintWriter w;
		String track = null;
		boolean foundTrack = false;
		
		try {
			s = new Scanner(new BufferedReader(new FileReader(input)));
			w = new PrintWriter(new BufferedWriter(new FileWriter(HMM_CONFIG_FILE)));
			
			if (posTagging) {
				posTagger = new PosTagger("left3words-wsj-0-18.tagger");
				File f = new File(POS_LOG_FILE);
				if (f.exists()) {
					Scanner log = new Scanner(new BufferedReader(new FileReader(POS_LOG_FILE)));
					track = log.next();
					log.close();
				}
			}
			else {
				testFiles = new Hashtable<String,String>();				
			}
			
			while (s.hasNext()) {
				String tag = s.next();
				String trainingFile = s.next();
				String testFile = s.next();
				
				String trainingOutFile;
				if (trainingFile.lastIndexOf('.') != -1) {
					trainingOutFile = TRAIN_FOLDER + trainingFile.substring(0, trainingFile.indexOf('.')) + ".seq";
				}
				else {
					trainingOutFile = TRAIN_FOLDER + trainingFile + ".seq";
				}
				
				String testOutFile;
				if (testFile.lastIndexOf('.') != -1) {
					testOutFile = TEST_FOLDER + testFile.substring(0, testFile.indexOf('.')) + ".seq";
				}
				else {
					testOutFile = TEST_FOLDER + testFile + ".seq";
				}

				if (posTagging) {
//					hmmThresholdModifiers.put(tag, 0.5);
//					System.out.println(hmmThresholdModifiers);
//					System.out.println(tag + " " + hmmThresholdModifiers.get(tag));
					posTagger.computeFrequentWords(true, STEMMER_FREQ_FOLDER, 0);
					posTagger.computeFrequentWords(false, STEMMER_FREQ_FOLDER, hmmThresholdModifiers.get(tag));
//					posTagger.computeFrequentWords(false, STEMMER_FREQ_FOLDER, 0.1);
					if (track != null && tag.compareTo(track) == 0) {
						foundTrack = true;
						continue;
					}
					
					if ((track != null && foundTrack) || (track == null && !foundTrack)) {
						PrintWriter log = new PrintWriter(POS_LOG_FILE);
						log.println(tag);
						log.close();
						File test = new File(POS_FOLDER + trainingFile);
//						if (!test.exists()) {
							posTagger.posTagUtterances(TEXT_FOLDER + trainingFile, POS_FOLDER + trainingFile, tag);
							posTagger.posTagUtterances(TEXT_FOLDER + testFile, POS_FOLDER + testFile, tag);
//						}
					}
				}
				else {
					testFiles.put(tag, testOutFile);
					
					Parser p = new Parser();
					int avg = p.encodePOSFile(dictionary, POS_FOLDER + trainingFile, trainingOutFile);
					p.encodePOSFile(dictionary, POS_FOLDER + testFile, testOutFile);
					
					// write data to the config file
					w.println(tag + " " + avg + " " + HMM_DELTA + " " + trainingOutFile);
				}
			}
			
			w.close();
			s.close();
			
			File f = new File(POS_LOG_FILE);
			f.delete();
		
		} catch (IOException e) {
			e.printStackTrace();
		}
	}
	
	public void train() {
		hmmSystem.train();
		nbc.train();
	}
	
	private void printSummary(PrintWriter pw, Hashtable<String,ArrayList<Integer>> summary) {
		for (String tag : summary.keySet()) {
			ArrayList<Integer> stats = summary.get(tag);
			pw.println(tag + " match rate: " + stats.get(0) + "% [" + stats.get(1) + "/" + stats.get(2) + "]");
		}
	}
	
	public Hashtable<String,ArrayList<Integer>> testHMM(String test_header) {
		System.out.println("\nTesting HMM system:");
		
		Hashtable<String,ArrayList<Integer>> summary = new Hashtable<String,ArrayList<Integer>>();
		
		try {
			PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter("logs/hmm/" + System.currentTimeMillis() + HMM_LOG_FILE)));
			pw.println(test_header);
			
			double matchRateSum = 0;
			for (String tag : testFiles.keySet()) {
				matchRateSum += hmmSystem.testTagFile(tag, testFiles.get(tag), pw, summary);
			}
			
			if (testFiles.size() > 0) {
				printSummary(pw, summary);
				System.out.println("Avg match rate: " + ((double) matchRateSum / testFiles.size()));
				pw.println("Avg match rate: " + ((double) matchRateSum / testFiles.size()));
			}
			
			pw.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
		
		return summary;
	}
	
	public void testNBC(String test_header) {
		System.out.println("\nTesting NBC model:");
		
		Hashtable<String,ArrayList<Integer>> summary = new Hashtable<String,ArrayList<Integer>>();
		
		try {
			PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter("logs/nbc/" + System.currentTimeMillis() + NBC_LOG_FILE)));
			pw.println(test_header);
			
			double matchRateSum = 0;
			for (String tag : testFiles.keySet()) {
				matchRateSum += nbc.testTagFile(tag, "text/test" + tag, pw, summary);
			}
			
			if (testFiles.size() > 0) {
				printSummary(pw, summary);
				System.out.println("Avg match rate: " + ((double) matchRateSum / testFiles.size()));
				pw.println("Avg match rate: " + ((double) matchRateSum / testFiles.size()));
			}
			
			pw.close();
		} catch (IOException e) {
			e.printStackTrace();
		}		
	}
	
	private String classifyUtterance(String utterance, List<ObservationInteger> seq) {
		StringTokenizer st = new StringTokenizer(utterance);
		
		if (st.countTokens() > HMM_THRESHOLD) {
			return hmmSystem.classifyUtterance(seq);
		}
		
		return nbc.classifyUtterance(utterance);
	}
	
	private int testTagFile(String tag, PrintWriter pw, Hashtable<String,ArrayList<Integer>> summary) {
		Reader reader;
		int match = 0;
		int uttCnt = 0;
		int matchRate = -1;
		
		ArrayList<Integer> track = new ArrayList<Integer>();
		
		try {
			// read HMM test sequences
			reader = new FileReader(testFiles.get(tag));
			List<List<ObservationInteger>> sequences = ObservationSequencesReader.readSequences(new ObservationIntegerReader(), reader);
			reader.close();

			// open utterances file
			BufferedReader bf = new BufferedReader(new FileReader("text/test" + tag));
			
			for (List<ObservationInteger> seq : sequences) {
				String utterance = bf.readLine();
				String bestMatch = classifyUtterance(utterance, seq);
				if (tag.compareTo(bestMatch) == 0) match ++;
				uttCnt ++;
				pw.println(utterance + " " + bestMatch + " " + tag);
			}
			
			bf.close();
			
			if (uttCnt == 0) return matchRate;
			matchRate = (match * 100) / uttCnt;
			track.add(matchRate);
			track.add(match);
			track.add(uttCnt);
			System.out.println(tag + " match rate: " + matchRate + "% [" + match + "/" + uttCnt + "]");
			pw.println(tag + " match rate: " + matchRate + "% [" + match + "/" + uttCnt + "]\n\n");
			
		} catch (Exception e) {
			e.printStackTrace();
		}
		
		return matchRate;
	}
	
	public double test(String test_header) {
		System.out.println("\nTesting system:");
		
		Hashtable<String,ArrayList<Integer>> summary = new Hashtable<String,ArrayList<Integer>>();
		
		try {
			PrintWriter pw = new PrintWriter(new BufferedWriter(new FileWriter("logs/" + System.currentTimeMillis() + LOG_FILE)));
			pw.println(test_header);
			
			int matchRateSum = 0;
			for (String tag : testFiles.keySet()) {
				matchRateSum += testTagFile(tag, pw, summary);
			}
			
			if (testFiles.size() > 0) {
				printSummary(pw, summary);
				System.out.println("Avg match rate: " + ((double) matchRateSum / testFiles.size()));
				pw.println("Avg match rate: " + ((double) matchRateSum / testFiles.size()));
				pw.close();	
				
				return (double) matchRateSum / testFiles.size();
			}
			
			pw.close();	
		} catch (IOException e) {
			e.printStackTrace();
		}
		
		return -1;
	}
}
